#pragma once

#include <queue>
#include <vector>

namespace solution_023
{

struct ListNode
{
    int val;
    ListNode* next;
    ListNode() : val(0), next(nullptr) {}
    ListNode(int x) : val(x), next(nullptr) {}
    ListNode(int x, ListNode* next) : val(x), next(next) {}
};

class Solution
{
public:
    ListNode* mergeKLists(std::vector<ListNode*>& lists)
    {
        if (lists.empty())
            return nullptr;

        auto gt = [&lists](int lhs, int rhs) { return lists.at(lhs)->val > lists.at(rhs)->val; };
        std::priority_queue<int, std::vector<int>, decltype(gt)> min_heap(gt);

        for (int i = 0; i < (int)lists.size(); ++i) {
            if (lists.at(i))
                min_heap.push(i);
        }

        ListNode root;
        for (auto prev = &root; !min_heap.empty(); ) {
            auto top = min_heap.top();
            auto node = new ListNode(lists.at(top)->val);
            prev->next = node;
            prev = node;
            min_heap.pop();

            lists.at(top) = lists.at(top)->next;
            if (lists.at(top))
                min_heap.push(top);
        }
        return root.next;
    }
};

}  /* namespace solution_023 */
